package main

import (
	"context"
	"fmt"
	"github.com/tx7do/kratos-transport/broker"
	"github.com/tx7do/kratos-transport/broker/kafka"
	"hash"
	"hash/fnv"
	"time"
)

const (
	// TODO 最好做成配置项！
	topic = "my-topic"
	//numPartitions = 6
)

func main() {

	userIds := []string{"user1", "user1", "user1", "user1", "user2", "user3", "user4"}
	fmt.Println("userIds: ", userIds)

	// 创建 Kafka 生产者
	ctx := context.Background()
	producer := kafka.NewBroker(
		broker.WithOptionContext(ctx),
		broker.WithAddress([]string{"localhost:9092"}...),
		broker.WithCodec("json"),
		//发送批次大小
		kafka.WithBatchSize(1),
		// 异步发送消息
		kafka.WithAsync(true),
	)

	err := producer.Init()
	if err != nil {
		panic(err)
	}
	err = producer.Connect()
	if err != nil {
		panic(err)
	}

	// 发送消息
	for idx, userId := range userIds {
		currStr := fmt.Sprintf("hello-%v-%v", idx, userId)
		msg := &broker.Message{
			Headers: broker.Headers{
				"userId": userId,
			},
			Body: currStr,
		}
		err = producer.Publish(
			topic,
			msg,
			// Notice 负载均衡器～根据用户id的hash指定分区发送消息
			kafka.WithHashBalancer(hashFunc(userId)),
			// Notice 同时需要将用户id使用WithMessageKey注入进去
			kafka.WithMessageKey([]byte(userId)),
		)
		fmt.Println("err: ", err)
		if err != nil {
			panic(err)
		}
	}

	// Notice 防止消息没发完 主协程停了 消息发不出去
	time.Sleep(time.Second * 2)

	producer.Disconnect()
}

func hashFunc(s string) hash.Hash32 {
	h := fnv.New32a()
	h.Write([]byte(s))
	return h
}
